# build_samples.py
# usage: 标记确认的check文件生成samples文件
# python ./src/tools/build_samples.py --path "./data" --samples "./data/samples.hand_coded"

import os
import json
import argparse
import glob

def main():
    parser = argparse.ArgumentParser(description="从指定路径下扫描 check 开头的文件，提取 flag != 0 的记录，写入 samples 文件。")
    parser.add_argument('--path', type=str, required=True, help='要扫描的目录路径')
    parser.add_argument('--samples', type=str, required=True, help='输出 samples 文件的完整路径（含文件名）')

    args = parser.parse_args()

    # 检查 path 是否存在
    if not os.path.isdir(args.path):
        print(f"错误：路径 '{args.path}' 不存在或不是目录。")
        return

    # 查找所有 check 开头的文件
    pattern = os.path.join(args.path, 'check*')
    files = glob.glob(pattern)
    if not files:
        print(f"警告：在路径 '{args.path}' 下未找到以 'check' 开头的文件。")
        return

    # 准备输出目录
    output_dir = os.path.dirname(args.samples)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 存储符合条件的样本
    valid_samples = []

    for file_path in files:
        if not os.path.isfile(file_path):
            continue
        print(f"处理文件: {file_path}")
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        data = json.loads(line)
                        # 只保留 flag != 0 的记录
                        if data.get('flag') != 0:
                            # 构造 samples 格式
                            sample = {
                                'id': data['id'],
                                'question': data['question'],
                                'sql': data['sql'],
                                'level': data['level']
                            }
                            valid_samples.append(sample)
                    except json.JSONDecodeError as e:
                        print(f"文件 {file_path} 第 {line_num} 行 JSON 解析错误: {e}")
                    except KeyError as e:
                        print(f"文件 {file_path} 第 {line_num} 行缺少字段: {e}")
        except Exception as e:
            print(f"读取文件 {file_path} 时出错: {e}")

    # 写入输出文件
    try:
        with open(args.samples, 'w', encoding='utf-8') as f:
            for sample in valid_samples:
                f.write(json.dumps(sample, ensure_ascii=False) + '\n')
        print(f"成功写入 {len(valid_samples)} 条记录到 '{args.samples}'")
    except Exception as e:
        print(f"写入文件 '{args.samples}' 失败: {e}")

if __name__ == '__main__':
    main()